#include "globalhomographycalculator.h"
#include <boost/foreach.hpp>
#include <map>
#include <queue>
#include "imagematch.h"
#include "inputimage.h"
#include "stitcher.h"
#include "imageset.h"

using namespace boost;
using namespace std;
using namespace cv;
using namespace cshotel;

GlobalHomographyCalculator::GlobalHomographyCalculator(shared_ptr<ImageSet> imgs) : imageSet(imgs) {}

void GlobalHomographyCalculator::calculate() {
    dout("Calculating global homographies ..." << endl);
    
    // first we find the center of the graph
    findCenter();
    dout("    center is " << center->getId() << endl);
    
    // then we do another Dijkstra from the center,
    // this time calculating the global homographies
    dout("    iterating pairwise homographies ..." << endl);
    assignHomographiesFrom(center);
}

double GlobalHomographyCalculator::assignHomographiesFrom(shared_ptr<InputImage> root) {
    map<shared_ptr<InputImage>, double> distance;
    multimap<double, shared_ptr<InputImage> > heap;
    map<shared_ptr<InputImage>, bool> visited;
    root->setMatchingVersionHomography(Mat::eye(3, 3, CV_64F));
    distance[root] = calculateRootDistance(root);
    double currentMaxDistance = distance[root];
    heap.insert(pair<double, shared_ptr<InputImage> >(distance[root], root));
    int numVisited = 0;
    while(numVisited < imageSet->getSize()) {
        shared_ptr<InputImage> current = heap.begin()->second;
        heap.erase(heap.begin());
        if(visited[current]) {
            continue;
        }
        visited[current] = true;
        ++numVisited;
        currentMaxDistance = max(currentMaxDistance, distance[current]);
        
        foreach(shared_ptr<ImageMatch> match, current->getImageMatches()) {
            if(shared_ptr<InputImage> firstImage = (match->getFirstImage().lock())) {
                if(shared_ptr<InputImage> secondImage = (match->getSecondImage().lock())) {
                    shared_ptr<InputImage> next;
                    if(firstImage == current) {
                        next = secondImage;
                    } else {
                        next = firstImage;
                    }
                    if(visited.find(next) != visited.end() && visited[next]) {
                        continue;
                    }
                    double nextDst = calculateNextDistance(distance[current], match);
                    if(distance.find(next) == distance.end() || nextDst < distance[next]) {
                        distance[next] = nextDst;
                        heap.insert(pair<double, shared_ptr<InputImage> >(distance[next], next));
                        next->setMatchingVersionHomography(current->getMatchingVersionHomography() * (Mat(match->getTransformationMatrix()).inv()));
                    }
                }
            }
        }
    }
}

void GlobalHomographyCalculator::findCenter() {
    dout("    finding center ..." << endl);
    
    double area = 1e10;
    center = *(imageSet->getImages().begin());
    foreach(shared_ptr<InputImage> img, imageSet->getImages()) {
        assignHomographiesFrom(img);
        ImageSet::BoundingBox boundingBox = imageSet->getBoundingBox();
        dout("        image " << img->getId() << " has bounding box (" << boundingBox.width << ", " << boundingBox.height << "), area = " << boundingBox.getArea() << endl);
        if(boundingBox.getArea() < area) {
            center = img;
            area = boundingBox.getArea();
        }
    }
}

// this distance function determines the order in which the nodes are visited
// by taking repeated maxima, we select the path that has the highest minimum
// number of RANSAC inliers
double GlobalHomographyCalculator::calculateNextDistance(double currentDistance, shared_ptr<ImageMatch> edge) {
    //return currentDistance + 1;
    return max(currentDistance, -static_cast<double>(edge->getInliers().size()));
}

double GlobalHomographyCalculator::calculateRootDistance(boost::shared_ptr<InputImage> root) {
    return -1000000000;
}
